import pandas as pd
import datetime
import math
import os
##-----------feature engineering---------------------

FD = 8
WS = 6
SP = 3
def set_rain_value(rain):
    if rain >= 999990 or rain == 0:
        return 0
    v = math.log(rain) + 2.5
    if v < 1:
        return 0.33
    elif v < 3:
        return 0.5
    elif v < 6:
        return 0.8
    else:
        return 1

def date_to_nth_day(date, format='%Y%m%d'):
    dd = datetime.datetime.strptime(str(int(date)), format)
    dd_struct = dd.timetuple()
    return dd_struct.tm_year, dd_struct.tm_mon, dd_struct.tm_yday


def fill_date_info(df):
    series_dates = df['date'].apply(date_to_nth_day)
    ser_year = series_dates.apply(lambda x:x[0])
    ser_days = series_dates.apply(lambda x:x[2])
    ser_month = series_dates.apply(lambda x: x[1])
    df['year'], df['days'], df['month'], df['weeks'] = \
        ser_year, ser_days,ser_month, ser_days//7


def get_avg_nfields(avg_fields, from_days , window_size , span ):
    return  ["{}{}{}".format(sf, rstart, window_size) for sf in avg_fields for rstart in range(from_days, (28 - window_size), span) ]


def get_avg_perfield(field, row, rs, ws):
    toavg = [row["{}{}".format(field, d)] for d in range(rs, rs+ws, 1) if row["{}{}".format(field, d)] < 999990]
    if len(toavg) == 0:
        print(field, rs, row['date'], row['station'], " substitution:({}-{})".format(max(rs-ws, 8), min(28, rs+ws+ws)))
        toavg = [row["{}{}".format(field, d)] for d in range(max(rs-ws, 8), min(28, rs+ws+ws), 1) if row["{}{}".format(field, d)] < 999990]
    return sum(toavg)/len(toavg)


def fill_pass_avg(df, shift_fields, from_days = FD, windows_size = WS, span = SP):
    ranges = range(from_days, (28 - windows_size), span)

    for rstart in ranges:
        for sf in shift_fields:
            df["{}{}{}".format(sf, rstart, windows_size)] = df.apply(lambda row: get_avg_perfield(sf, row, rstart, windows_size), axis=1)
    return get_avg_nfields(shift_fields, from_days, windows_size, span)



def feature_engineering(df, reload = True, for_train =True):
    tmp_df = "tmp.csv"

    if reload and os.path.exists(tmp_df):
        print("load feature_engineering old")
        df = pd.read_csv(tmp_df)
    else:
        print("start feature_engineering")
        fill_date_info(df) #parse date time
        def city_month(city, month):
            return city + str(month)
        df['cm'] = df.apply(lambda row: city_month(row['city'], row['month']), axis=1)

        #collect passed info
        pass_fields = fill_pass_avg(df, shift_fields= ['temperature', 'humidity', 'cloud', 'rain', 'wind_speed'])
        if for_train == True:
            print("save for next time")
            df[['year','month','city', 'station', 'cm', 'days','temperature', 'humidity', 'rain', 'wind_speed'] + pass_fields].to_csv(tmp_df, index=False)
    for rf in df.columns:
        if rf.startswith('rain'):
            df[rf] = df[rf].apply(set_rain_value)
    return df
#======================================



#=====preprocessing================================================
def fill_missing_within_city(df, null_fields = ['rain20', 'wind_speed','rain08' ,'cloud','humidity', 'temperature', 'pressure']):
    grp = df.groupby(['city', 'date'])
    print_flag = 0
    for k, gdata in grp:
        if int(k[1]/10000) != print_flag:
            print("current:{}".format(k))
            print_flag = int(k[1]/10000)

        for null_fld in null_fields:
            missing = gdata[(gdata[null_fld].isnull()) | (gdata[null_fld] >= 999990)]
            if len(missing) > 0:
                corrects = gdata[ (gdata[null_fld] < 999990)] # no need to check NaN
                if len(corrects) > 0:
                    avg_of_city = corrects[null_fld].mean()
                    df.loc[missing.index, null_fld] = 0 if null_fld.startswith("rain") and avg_of_city < 0.083 else avg_of_city



def shift_data(df,  outpath, fields = ['temperature', 'humidity', 'rain', 'cloud', 'wind_speed'], days = 3):
    grp = df.groupby(['station'])
    mode = 'w'
    header = True
    for k, gdata in grp:
        v = gdata.sort_values(['date'])
        for field in fields:
            for d in range(1, days + 1):
                newcol= v[field].shift(d)
                for pd in range(0,d):
                    newcol.iloc[pd]= 0
                v["{}{}".format(field,d)] = newcol
        v.iloc[days:].to_csv(outpath, mode=mode, header=header, index=False)
        mode = 'a'
        header = False
        print(k)


def fill_and_shift(df, outname, null_fields = ['rain20', 'wind_speed','rain08' ,'cloud','humidity', 'temperature', 'pressure'],
                   shift_fields =['temperature', 'humidity', 'rain', 'cloud', 'pressure','wind_speed'], shift_days = 2):
    print("fill missing first")
    fill_missing_within_city(df, null_fields)

    df['rain'] = (df['rain20'] + df['rain08'])/2
    df.to_csv(outname + ".fil", index=False)
    shift_data(df, outname + ".csv", fields= shift_fields, days= shift_days)


def split_original(df, train_out = 'train0.csv', test_out = 'test0.csv'):
    df_train = df[df['date']<= 20171201]
    df_test = df[df['date'] > 20171101]
    df_train.to_csv(train_out, index=False)
    df_test.to_csv(test_out, index=False)



def get_next7days(date = "20171231",  full_columns = ['date', 'station'], station = 100):
    dd = datetime.datetime.strptime(str(int(date)), '%Y%m%d')
    df = []
    for d in range(1,8):
        dn = (dd + datetime.timedelta(days= d)).timetuple()
        df.append([dn.tm_year* 10000 +  dn.tm_mon * 100 + dn.tm_mday,station ] + [0 for x in full_columns[2:]] )
    return pd.DataFrame(df, columns=  full_columns)


def gen_test_submit(df_test, outpath):
    dt_from = df_test['date'].max()
    fields = list(df_test.columns.values)
    mode = 'w'
    header = True
    for station, gdata in df_test.groupby(['station', 'city']):
        v = gdata.sort_values(['date'])
        df_submit = get_next7days(dt_from, fields, station = station[0])
        df_submit['city'] = station[1]
        v = v.append( df_submit , ignore_index=True)
        v.to_csv(outpath, mode=mode, header=header, index=False)
        mode = 'a'
        header = False
        print(station )
    return dt_from

def gen_only_submit(df_test, outpath):
    dt_from = df_test['date'].max()
    fields = list(df_test.columns.values)
    mode = 'w'
    header = True
    for station, gdata in df_test.groupby(['station', 'city']):
        df_submit = get_next7days(dt_from, fields, station = station[0])
        df_submit['city'] = station[1]
        df_submit.to_csv(outpath, mode=mode, header=header, index=False)
        mode = 'a'
        header = False
        print(station )
    return dt_from



if __name__ == "__main__":
    df = pd.read_csv("d:/sfxy191/bd2019-weather-prediction-training-20190608.csv")
    split_original(df, train_out='train0.csv', test_out='test0.csv')

    df.sort_values(['date','city', 'station']).to_csv("www.csv", index=False)

    df = pd.read_csv("tmp.csv")
    cols = set(df.columns)
    # cols.difference_update(['cm','city','month','year','rain','station'])
    corr = (df[list(cols)].corr(method='kendall'))
    print(corr['temperature'].sort_values())
    print("==========")
    print(corr['humidity'].sort_values())

  #  split_original(df, test_out='test0.csv')
 #   gen_test_submit(pd.read_csv('test0.csv'), 'test1.csv')


    #fill_missing_within_city(df)

    #df.to_csv(fill_miss, index=False)
  #  df = pd.read_csv(fill_miss)
  #  shift_data(df, "fill_miss_shift.csv")
    #stat_pass_to_now(df, "train.csv")
    # df = pd.read_csv("test.csv")
    # print("-----")
    # simple_flatten(df)
    # df.to_csv("train1.csv", index=False)